import typing as tp

from loguru import logger
import viztracer

from ..task import Task
from ..agent import Agent
from ..common import Result
from ..jobs.single_node_single_mod_job import Step
from . import type2executor
from ..events.signal import TaskEvent
from ..executors.base_aio_process import AIOCmdExecutor


class SNSMRunner:
    def __init__(self, agent: Agent, task: Task):
        self.agent = agent
        self.task = task
        self._is_interrupted = False

    async def run(self):
        try:
            await TaskEvent.start_signal.send_async(self.__class__.__name__, agent=self.agent, task=self.task)
            if not self.agent.is_unique_running:
                if self.task.job.is_unique:
                    self.agent.is_unique_running = True
                if self.task.job.enable_viztracer:
                    with viztracer.VizTracer(
                            log_async=True, include_files=["*/mnagent/*"], output_file=self.task.viztracer_result_json
                    ):
                        await self.run_step(self.task.job.first_step)
                else:
                    await self.run_step(self.task.job.first_step)
            else:
                self.task.is_delayed = True
            await TaskEvent.finish_signal.send_async(self.__class__.__name__, agent=self.agent, task=self.task)
        except (BaseException, Exception):
            import traceback
            logger.error(traceback.format_exc())
        finally:
            if self.task.job.is_unique and not self.task.is_delayed:
                self.agent.is_unique_running = False
            await TaskEvent.finally_signal.send_async(self.__class__.__name__, agent=self.agent, task=self.task)

    async def run_step(self, step: Step):
        await TaskEvent.step_start_signal.send_async(
            self.__class__.__name__, agent=self.agent, task=self.task, step=step)
        exe = type2executor.get(step.executor.type)
        if exe is None:
            result = Result.EXECUTOR_NOT_FOUND
            await TaskEvent.step_finish_signal.send_async(
                self.__class__.__name__, agent=self.agent, task=self.task, step=step, result=result)
        else:
            if not self._is_interrupted:
                async def _executor_log(message: str, type_: tp.Literal["STDERR", "STDOUT"]):
                    await TaskEvent.executor_log_signal.send_async(
                        self.__class__.__name__, agent=self.agent, task=self.task, message=message, type_=type_)

                exe = exe(step.executor)
                self.task.executors.add(exe)
                code = await exe.run(self.agent, log_callback=_executor_log)
                self.task.executors.remove(exe)

                if code == AIOCmdExecutor.CMD_ERROR_CODE:
                    result = Result.EXECUTOR_CMD_ERROR
                elif code == AIOCmdExecutor.TIMEOUT_CODE:
                    result = Result.EXECUTOR_TIMEOUT
                elif code == AIOCmdExecutor.INTERRUPT_CODE:
                    result = Result.EXECUTOR_INTERRUPT
                elif code == AIOCmdExecutor.UNKNOWN_CODE:
                    result = Result.EXECUTOR_UNKNOWN
                elif code in step.validate_code:
                    result = Result.SUCCESS
                else:
                    result = Result.FAIL

            else:
                result = Result.TASK_INTERRUPT

            await TaskEvent.step_finish_signal.send_async(
                self.__class__.__name__, agent=self.agent, task=self.task, step=step, result=result)

            if result.value == Result.TASK_INTERRUPT or (step.next and code in step.next.ready_validate_code):
                return await self.run_step(step.next.step)

    async def interrupt(self):
        self._is_interrupted = True
        for exe in self.task.executors:
            await exe.interrupt()
